#!/usr/bin/env python

import numpy as np
import math
from matplotlib.mlab import frange

def rectangular(nx,ny,m):
    """Migration array for populations on a rectangular grid
    with dimensions nx*ny.
    Migration between any two populations occurs symmetrically at
    a rate m.
    
    For ny=1, a linear stepping stone population structure results."""
    n = nx * ny
    migrates = np.zeros((n,n))
    
    # one population only
    if nx == ny == 1:
        migrates[0,0] = 1.
        return migrates
    
    # horizontal linear stepping stones of nx populations
    if ny == 1:
        for k in range(1,n-1):
            migrates[k,k] = 1-2*m
            migrates[k,k-1] = m
            migrates[k,k+1] = m
        migrates[0,0] = 1-m
        migrates[0,1] = m
        migrates[n-1,n-1] = 1-m
        migrates[n-1,n-2] = m
        return migrates
    
    # vertical linear stepping stones of ny populations
    if nx == 1:
        migrates = rectangular(ny,nx,m)
        return migrates.T
    
    # now for the "real" rectangles
    TLcorner = 0
    TRcorner = nx-1
    BLcorner = n-nx
    BRcorner = n-1
    Ledge = frange(nx,n-2*nx,nx)
    Redge = frange(2*nx-1,n-nx-1,nx)
    Tedge = frange(1,nx-2,1)
    Bedge = frange(n-nx+1,n-2,1)
#    print """
#    corners:
#        top left:     %d
#        top right:    %d
#        bottom left:  %d
#        bottom right: %d
#    edges:
#        left:   %s
#        right:  %s
#        top:    %s
#        bottom: %s""" % (TLcorner,
#                         TRcorner,
#                         BLcorner,
#                         BRcorner,
#                         str(Ledge),
#                         str(Redge),
#                         str(Tedge),
#                         str(Bedge))
    
    for k in range(n):
        if k == TLcorner:
            migrates[k,k] = 1-2*m
            for p in [k+1,k+nx]:
                migrates[k,p] = m
        elif k == TRcorner:
            migrates[k,k] = 1-2*m
            for p in [k-1,k+nx]:
                migrates[k,p] = m
        elif k == BLcorner:
            migrates[k,k] = 1-2*m
            for p in [k-nx,k+1]:
                migrates[k,p] = m
        elif k == BRcorner:
            migrates[k,k] = 1-2*m
            for p in [k-nx,k-1]:
                migrates[k,p] = m
        elif k in Ledge:
            migrates[k,k] = 1-3*m
            for p in [k-nx,k+1,k+nx]:
                migrates[k,p] = m
        elif k in Redge:
            migrates[k,k] = 1-3*m
            for p in [k-nx,k-1,k+nx]:
                migrates[k,p] = m 
        elif k in Tedge:
            migrates[k,k] = 1-3*m
            for p in [k-1,k+1,k+nx]:
                migrates[k,p] = m 
        elif k in Bedge:
            migrates[k,k] = 1-3*m
            for p in [k-nx,k-1,k+1]:
                migrates[k,p] = m  
        else:
            migrates[k,k] = 1-4*m
            for p in [k-nx,k-1,k+1,k+nx]:
                migrates[k,p] = m        
    return migrates


def square(n,m):
    """Migration array for n populations on a square grid."""
    nx = ny = int(math.sqrt(n))
    if not (math.sqrt(n) % nx == 0.):
        raise ValueError, "number of populations (n) must be the square of an integer"
    return rectangular(nx,ny,m)

def rectangular_hybrid_zone(nx,ny,m):
    """Migration array for the following population structure:
        Between two source populations there is a hybrid zone with
        populations on a rectangular grid.
        
        E.g.:
           1 - 2 - 3
        0  |   |   |  7
           4 - 5 - 6
           """
    # rectangular hybrid zone plus two source populations
    nHZ = nx * ny    # number of populations in the hybrid zone
    nT = nHZ + 2     # total number of populations
    migrates = np.zeros((nT,nT))
    
    TLcorner = 1
    TRcorner = nx
    BLcorner = nHZ-nx+1
    BRcorner = nHZ
    Ledge = frange(nx+1,nHZ-2*nx+1,nx)
    Redge = frange(2*nx,nHZ-nx,nx)
    Tedge = frange(2,nx-1,1)
    Bedge = frange(nHZ-nx+2,nHZ-1,1)
    
    migrates[1:nHZ+1,1:nHZ+1] = rectangular(nx,ny,m)
    migrates[0,0] = migrates[nT-1,nT-1] = 1.
    
    for k in range(1,nT-1):
        if ((k==TLcorner) or (k==BLcorner) or (k in Ledge)):
            migrates[k,k] -= m
            migrates[k,0] = m
        if ((k==TRcorner) or (k==BRcorner) or (k in Redge)):
            migrates[k,k] -= m
            migrates[k,nT-1] = m
    
    return migrates


if __name__ == '__main__':
    m = 0.01
    
    ns = 2
    mig2x2 = rectangular(ns,ns,m)
    print "migration on a %dx%d grid:" % (ns,ns)
    print mig2x2
        
    nx = 5
    ny = 3
    mig5x3 = rectangular(nx,ny,m)
    print "\nmigration on a %dx%d grid:" % (nx,ny)
    print mig5x3
    
    
    
